Skip to content

Conversation

@kingcrimsontianyu
Copy link
Contributor

@kingcrimsontianyu kingcrimsontianyu commented Jan 6, 2026

This PR exposes cuFile stream register and deregister functions to KvikIO Python API.

kvikio.stream_register(stream, flags)
kvikio.stream_deregister(stream)

The stream register and deregister is an optional, performance-tuning API. The users specify the precondition of the I/O parameters by setting the flags that indicate when the I/O parameters become valid and whether they are page aligned.

The added Python API can be unit tested using:

pytest ~/kvikio/python/kvikio/tests/test_async_io.py::test_stream_register_deregister -m cufile -vsx

A known issue is that, under CUDA 12.9, which is the only setting tested, if the file path is not on a drive with full GDS support, the above unit test will pass but end with a segmentation fault. This is beyond the scope of this PR and is likely not an issue in KvikIO. Further investigation is a good-to-have.

@kingcrimsontianyu kingcrimsontianyu added feature request New feature or request non-breaking Introduces a non-breaking change python Affects the Python API of KvikIO c++ Affects the C++ API of KvikIO labels Jan 6, 2026
@copy-pr-bot
Copy link

copy-pr-bot bot commented Jan 6, 2026

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@kingcrimsontianyu kingcrimsontianyu marked this pull request as ready for review January 8, 2026 17:31
@kingcrimsontianyu kingcrimsontianyu requested review from a team as code owners January 8, 2026 17:31
Copy link
Member

@madsbk madsbk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We might consider using something like the following for the CUDA Stream bindings

from rmm.pylibrmm.stream import Stream

https://docs.rapids.ai/api/rmm/stable/python/pylibrmm/#rmm.pylibrmm.stream.Stream

@kingcrimsontianyu
Copy link
Contributor Author

kingcrimsontianyu commented Jan 8, 2026

I've changed the parameter name to raw_stream for the Python API. For your suggestion, do you think we need to add Python RMM dependency to KvikIO, or add KvikIO's own stream wrapper on C++ and Python level? @madsbk

Copy link
Member

@madsbk madsbk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like the raw_stream name.

For your suggestion, do you think we need to add Python RMM dependency to KvikIO, or add KvikIO's own stream wrapper on C++ and Python level? @madsbk

At some point, yes. Or use cuda-python's CUDA stream, or roll our own :)

@kingcrimsontianyu
Copy link
Contributor Author

That makes sense. I've created this ticket #895 .
Re-requested the review after a minor update to the cufile.py which uses the raw stream as well.

@kingcrimsontianyu
Copy link
Contributor Author

/merge

@rapids-bot rapids-bot bot merged commit e7ffd20 into rapidsai:main Jan 13, 2026
84 checks passed
@GregoryKimball GregoryKimball moved this from Burndown to Landed in libcudf Jan 13, 2026
@kingcrimsontianyu
Copy link
Contributor Author

The segmentation fault mentioned in the PR description is actually caused by a tricky bug in Cython stream casting, fixed by #904.

rapids-bot bot pushed a commit that referenced this pull request Jan 20, 2026
#893 contains a tricky bug. The following is incorrect:

```python
def stream_register(stream: int, flags: int) -> None:
    cdef CUstream cpp_stream = <CUstream>stream
```
`stream` is a Python object. The cast applies to the address of this Python object, instead of the underlying integer value. This surprisingly does not trigger a compile error.

There are two ways to fix this bug:
- Method 1: double casting. The additional cast to `uintptr_t` extracts the underlying integer value of `stream` correctly.
```python
def stream_register(stream: int, flags: int) -> None:
    cdef CUstream cpp_stream = <CUstream><uintptr_t>stream
```

- Method 2: `uintptr_t` as type hint (used in this PR). `uintptr_t` is not just a type hint; it actually causes `stream` to be a C object.
```python
def stream_register(stream: uintptr_t, flags: int) -> None:
    cdef CUstream cpp_stream = <CUstream>stream
```
For both method, `uintptr_t` must be `cimport`-ed into the source file.

Authors:
  - Tianyu Liu (https://github.com/kingcrimsontianyu)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)

URL: #904
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

c++ Affects the C++ API of KvikIO feature request New feature or request non-breaking Introduces a non-breaking change python Affects the Python API of KvikIO

Projects

Status: Landed

Development

Successfully merging this pull request may close these issues.

3 participants